package com.zoe.onelink.logging.trace.dubbo;

import brave.Span;
import brave.Span.Kind;
import brave.propagation.CurrentTraceContext;
import brave.propagation.CurrentTraceContext.Scope;
import brave.propagation.TraceContext;
import brave.rpc.RpcClientHandler;
import brave.rpc.RpcServerHandler;
import brave.rpc.RpcTracing;
import cn.hutool.core.text.StrFormatter;
import com.alibaba.fastjson.JSON;
import com.zoe.onelink.common.constant.TraceSpanConstant;
import com.zoe.onelink.logging.trace.property.TraceLogProperties;
import com.zoe.onelink.logging.trace.util.TraceLogUtil;
import com.zoe.onelink.util.PathMatchUtil;
import com.zoe.onelink.util.SpringUtil;
import org.apache.dubbo.common.constants.CommonConstants;
import org.apache.dubbo.common.extension.Activate;
import org.apache.dubbo.rpc.Filter;
import org.apache.dubbo.rpc.Invocation;
import org.apache.dubbo.rpc.Invoker;
import org.apache.dubbo.rpc.Result;
import org.apache.dubbo.rpc.RpcContext;
import org.apache.dubbo.rpc.RpcException;

import java.util.Map;
import java.util.concurrent.CompletableFuture;

import static brave.internal.Throwables.propagateIfFatal;

/**
 * <p>标题: </p>
 * <p>描述: </p>
 * <p>版权: Copyright (c) 2021</p>
 * <p>公司: 智业软件股份有限公司</p>
 *
 * @version: 1.0
 * @author: ruansheng
 * @date: 2021-08-09
 */
@Activate(group = {CommonConstants.PROVIDER, CommonConstants.CONSUMER}, order = 200)
public final class DubboTraceFilter implements Filter {

    private TraceLogProperties.DubboLogProperties dubboLogProperties;
    CurrentTraceContext currentTraceContext;
    RpcClientHandler clientHandler;
    RpcServerHandler serverHandler;
    private volatile boolean isInit = false;

    /**
     * ExtensionLoader提供的跟踪实现必须被命名为“rpcTracing”。例如，如果使用SpringExtensionFactory，只会注入一个名为“rpcTracing”的bean。
     * 自定义解析
     * 自定义解析器，如RpcTracing.clientRequestParser()，可以使用特定于dubbo的类型DubboRequest和DubboResponse来获得诸如Java调用或结果之类的访问。
     */
    public void setRpcTracing(RpcTracing rpcTracing) {
        if (rpcTracing == null) {
            throw new NullPointerException("rpcTracing == null");
        }
        currentTraceContext = rpcTracing.tracing().currentTraceContext();
        clientHandler = RpcClientHandler.create(rpcTracing);
        serverHandler = RpcServerHandler.create(rpcTracing);
        isInit = true;
    }

    public void setTraceLogProperties(TraceLogProperties traceLogProperties) {
        if (traceLogProperties != null) {
            this.dubboLogProperties = traceLogProperties.getDubbo();
        }
    }

    @Override
    public Result invoke(Invoker<?> invoker, Invocation invocation) throws RpcException {
        if (this.dubboLogProperties == null) {
            TraceLogProperties traceLogProperties = SpringUtil.getBean(TraceLogProperties.class);
            if (traceLogProperties != null) {
                this.setTraceLogProperties(traceLogProperties);
            }
        }
        if (!isInit || !this.dubboLogProperties.isEnabled()) {
            return invoker.invoke(invocation);
        }

        TraceContext invocationContext = currentTraceContext.get();

        RpcContext rpcContext = RpcContext.getContext();
        Kind kind = rpcContext.isProviderSide() ? Kind.SERVER : Kind.CLIENT;
        Span span;
        DubboRequest request;
        if (kind.equals(Kind.CLIENT)) {
            // 当一个服务调用B服务，然后B服务再调用C服务时
            // C服务跨度是A当从调用. getattachments()读取。这是因为
            // AbstractInvoker通过RpcContext.getContext()添加附件，而不是调用。
            // org.apache.dubbo.rpc.protocol。AbstractInvoker(第141行)来自v2.7.3
            Map<String, String> attachments = RpcContext.getContext().getAttachments();
            DubboClientRequest clientRequest = new DubboClientRequest(invoker, invocation, attachments);
            request = clientRequest;
            span = clientHandler.handleSendWithParent(clientRequest, invocationContext);

            String invokeMethodName = StrFormatter.format("{}.{}", invoker.getInterface().getName(), invocation.getMethodName());
            boolean isAllowMethodParamLog = PathMatchUtil.isPermit(invokeMethodName, this.dubboLogProperties.getIncludeMethods(),
                    this.dubboLogProperties.getExcludeMethods(), true, true, "DUBBO_TRACE");
            // 记录Dubbo入参: 1.当前为Consumer 2.Dubbo日志已启用 3.方法路径匹配成功
            if (this.dubboLogProperties.isEnabled() && isAllowMethodParamLog) {
                String dubboLog = TraceLogUtil.trim(JSON.toJSONString(invocation.getArguments()), dubboLogProperties.getTrimType(), dubboLogProperties.getParamMaxLength());
                span.tag(TraceSpanConstant.RPC_PARAM, dubboLog);
            }
        } else {
            DubboServerRequest serverRequest = new DubboServerRequest(invoker, invocation);
            request = serverRequest;
            span = serverHandler.handleReceive(serverRequest);
        }
        span.tag(TraceSpanConstant.LOG_TYPE, TraceSpanConstant.LogType.DUBBO);

        boolean isSynchronous = true;
        Scope scope = currentTraceContext.newScope(span.context());
        Result result = null;
        Throwable error = null;
        try {
            result = invoker.invoke(invocation);
            error = result.getException();
            CompletableFuture<Object> future = rpcContext.getCompletableFuture();
            if (future != null) {
                isSynchronous = false;
                // 注意: 目前没有检测CompletableFuture，所以回调将不会看到
                // 调用上下文，除非它们使用由CurrentTraceContext检测的执行器
                // 如果稍后测试这个，注意使用正确的上下文依赖于RPC类型!
                future.whenComplete(FinishSpan.create(this, request, result, span));
            }
            return result;
        } catch (Throwable e) {
            propagateIfFatal(e);
            error = e;
            throw e;
        } finally {
            if (isSynchronous) {
                FinishSpan.finish(this, request, result, error, span);
            }
            scope.close();
        }
    }
}
